import os

import insightface
import cv2
import numpy as np
import random
import uuid
import tqdm

class PersonDet():

    def __init__(self):
        self.bdet = insightface.model_zoo.get_model('scrfd_person_2.5g.onnx', download=True)
        self.bdet.prepare(0, det_thresh=0.75, nms_thresh=0.6, input_size=(640, 640))


    def _body_detect(self, image: np.ndarray):
        bboxes, kpss = self.bdet.detect(image)
        return self.detect_person(image, bboxes, kpss)


    @staticmethod
    def detect_person(img: np.ndarray, bboxes, kpss):
        bboxes = np.round(bboxes[:, :4]).astype(np.int)
        kpss = np.round(kpss).astype(np.int)
        kpss[:, :, 0] = np.clip(kpss[:, :, 0], 0, img.shape[1])
        kpss[:, :, 1] = np.clip(kpss[:, :, 1], 0, img.shape[0])
        vbboxes = bboxes.copy()
        vbboxes[:, 0] = kpss[:, 0, 0]
        vbboxes[:, 1] = kpss[:, 0, 1]
        vbboxes[:, 2] = kpss[:, 4, 0]
        vbboxes[:, 3] = kpss[:, 4, 1]
        return bboxes, vbboxes

    @staticmethod
    def draw(img: np.ndarray, bboxes, vbboxes):
        for i in range(bboxes.shape[0]):
            bbox = bboxes[i]
            vbbox = vbboxes[i]
            x1, y1, x2, y2 = bbox
            vx1, vy1, vx2, vy2 = vbbox
            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 1)
            alpha = 0.8
            color = (255, 0, 0)
            for c in range(3):
                img[vy1:vy2, vx1:vx2, c] = img[vy1:vy2, vx1:vx2, c] * alpha + color[c] * (1.0 - alpha)
            cv2.circle(img, (vx1, vy1), 1, color, 2)
            cv2.circle(img, (vx1, vy2), 1, color, 2)
            cv2.circle(img, (vx2, vy1), 1, color, 2)
            cv2.circle(img, (vx2, vy2), 1, color, 2)

def apply_augmentations(bbox, img_shape, max_translation, max_rotation, max_scale):
    x1, y1, x2, y2 = bbox
    center_x, center_y = (x1 + x2) / 2, (y1 + y2) / 2

    # 平移
    translation_x = random.randint(-max_translation, max_translation)
    translation_y = random.randint(-max_translation, max_translation)

    # 旋转和缩放
    angle = random.uniform(-max_rotation, max_rotation)
    scale = random.uniform(1 - max_scale, 1 + max_scale)
    M = cv2.getRotationMatrix2D((float(center_x), float(center_y)), angle, scale)

    # 应用变换
    points = np.array([[x1, y1], [x2, y1], [x1, y2], [x2, y2]])
    ones = np.ones(shape=(len(points), 1))
    points_ones = np.hstack([points, ones])

    # 变换所有点
    transformed_points = M.dot(points_ones.T).T

    # 计算新的边界框
    x1_new, y1_new = np.min(transformed_points, axis=0)[:2]
    x2_new, y2_new = np.max(transformed_points, axis=0)[:2]

    x1_new = max(0, int(x1_new + translation_x))
    y1_new = max(0, int(y1_new + translation_y))
    x2_new = min(img_shape[1], int(x2_new + translation_x))
    y2_new = min(img_shape[0], int(y2_new + translation_y))

    return [x1_new, y1_new, x2_new, y2_new]


def process_images(images_folder, save_folder, max_translation, max_rotation, max_scale, debug, pdet):
    # 确保保存文件夹存在
    if not debug and not os.path.exists(save_folder):
        os.makedirs(save_folder)

    # 遍历文件夹中的图像
    for filename in tqdm.tqdm(os.listdir(images_folder)):
        if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
            file_path = os.path.join(images_folder, filename)
            img = cv2.imread(file_path)
            img_shape = img.shape

            # 使用PersonDet进行行人检测
            bboxes, vbboxes = pdet._body_detect(image=img)

            for i, bbox in enumerate(bboxes):
                # 应用随机扰动
                augmented_bbox = apply_augmentations(bbox, img_shape, max_translation, max_rotation, max_scale)
                bboxes[i] = augmented_bbox

                if not debug:
                    # 裁剪并保存图像
                    x1, y1, x2, y2 = augmented_bbox
                    cropped_img = img[y1:y2, x1:x2]
                    save_path = os.path.join(save_folder, str(uuid.uuid4()) + ".jpg")
                    cv2.imwrite(save_path, cropped_img)

            if debug:
                # 在debug模式下，只显示图像
                PersonDet.draw(img, bboxes, vbboxes)
                cv2.imshow("Detected Persons", img)
                cv2.waitKey(0)
                cv2.destroyAllWindows()


pdet = PersonDet()

# 示例使用
images_folder = "/Users/tunm/Downloads/val2017"
save_folder = "/Users/tunm/Downloads/val2017_other"
max_translation = 5  # 最大平移像素
max_rotation = 5     # 最大旋转角度
max_scale = 0.1       # 最大缩放比例
debug = False          # 或者 False

# 调用函数
process_images(images_folder, save_folder, max_translation, max_rotation, max_scale, debug, pdet)